# File with code needed to fit ordered beta regression
# As custom family in R package brms
# See vignette at 

# custom family

ord_beta_reg <- custom_family("ord_beta_reg",
                              dpars=c("mu","phi","cutzero","cutone"),
                              links=c("logit","log",NA,NA),
                              lb=c(NA,0,NA,NA),
                              type="real")

# stan code for density of the model

stan_funs <- "real ord_beta_reg_lpdf(real y, real mu, real phi, real cutzero, real cutone) {

    //auxiliary variables
    real mu_logit = logit(mu);
    vector[2] thresh;
    thresh[1] = cutzero;
    thresh[2] = cutzero + exp(cutone);
    
  if(y==0) {
      return log1m_inv_logit(mu_logit - thresh[1]);
    } else if(y==1) {
      return log_inv_logit(mu_logit  - thresh[2]);
    } else {
      return log(inv_logit(mu_logit  - thresh[1]) - inv_logit(mu_logit - thresh[2])) +
                beta_proportion_lpdf(y|mu,phi);
    }
  }"

stanvars <- stanvar(scode=stan_funs,block="functions")

# For pulling posterior predictions

posterior_predict_ord_beta_reg <- function(i, draws, ...) {
  mu <- draws$dpars$mu[, i]
  phi <- draws$dpars$phi
  cutzero <- draws$dpars$cutzero
  cutone <- draws$dpars$cutone
  N <- length(phi)
  
  thresh1 <- cutzero
  thresh2 <- cutzero + exp(cutone)
  
  pr_y0 <- 1 - plogis(qlogis(mu) - thresh1)
  pr_y1 <- plogis(qlogis(mu) - thresh2)
  pr_cont <- plogis(qlogis(mu)-thresh1) - plogis(qlogis(mu) - thresh2)
  out_beta <- rbeta(n=N,mu*phi,(1-mu)*phi)
  
  # now determine which one we get for each observation
  outcomes <- sapply(1:N, function(i) {
    sample(1:3,size=1,prob=c(pr_y0[i],pr_cont[i],pr_y1[i]))
  })
  
  final_out <- sapply(1:length(outcomes),function(i) {
    if(outcomes[i]==1) {
      return(0)
    } else if(outcomes[i]==2) {
      return(out_beta[i])
    } else {
      return(1)
    }
  })
  
  final_out
  
}

# for calculating marginal effects/conditional expectations

posterior_epred_ord_beta_reg<- function(draws) {
  
  cutzero <- draws$dpars$cutzero
  cutone <- draws$dpars$cutone
  
  mu <- draws$dpars$mu
  
  thresh1 <- cutzero
  thresh2 <- cutzero + exp(cutone)
  
  low <- 1 - plogis(qlogis(mu) - thresh1)
  middle <- plogis(qlogis(mu)-thresh1) - plogis(qlogis(mu) - thresh2)
  high <- plogis(qlogis(mu) - thresh2)
  
  low*0 + middle*mu + high
}

# for calcuating LOO and Bayes Factors

log_lik_ord_beta_reg <- function(i, draws) {
  
  mu <- draws$dpars$mu[,i]
  phi <- draws$dpars$phi
  y <- draws$data$Y[i]
  cutzero <- draws$dpars$cutzero
  cutone <- draws$dpars$cutone
  
  thresh1 <- cutzero
  thresh2 <- cutzero + exp(cutone)
  
  if(y==0) {
    out <- log(1 - plogis(qlogis(mu) - thresh1))
  } else if(y==1) {
    out <- log(plogis(qlogis(mu) - thresh2))
  } else {
    out <- log(plogis(qlogis(mu)-thresh1) - plogis(qlogis(mu) - thresh2)) + dbeta(y,mu*phi,(1-mu)*phi,log=T)
  }
  
  out
  
}

###### Code declaring induced dirichlet prior ####
# code from Michael Betancourt/Staffan Betner
# discussion here: https://discourse.mc-stan.org/t/dirichlet-prior-on-ordinal-regression-cutpoints-in-brms/20640
dirichlet_prior <- "
  real induced_dirichlet_lpdf(vector c, vector alpha, real phi) {
    int K = num_elements(c) + 1;
    vector[K - 1] sigma = inv_logit(phi - c);
    vector[K] p;
    matrix[K, K] J = rep_matrix(0, K, K);
    
    // Induced ordinal probabilities
    p[1] = 1 - sigma[1];
    for (k in 2:(K - 1))
      p[k] = sigma[k - 1] - sigma[k];
    p[K] = sigma[K - 1];
    
    // Baseline column of Jacobian
    for (k in 1:K) J[k, 1] = 1;
    
    // Diagonal entries of Jacobian
    for (k in 2:K) {
      real rho = sigma[k - 1] * (1 - sigma[k - 1]);
      J[k, k] = - rho;
      J[k - 1, k] = rho;
    }
    
    return   dirichlet_lpdf(p | alpha)
           + log_determinant(J);
  }
"
dirichlet_prior_stanvar <- stanvar(scode = dirichlet_prior, block = "functions")

stanvar(scode = "ordered[2] thresh;
              thresh[1] = cutzero;
              thresh[2] = cutzero+exp(cutone);", 
        block = "tparameters") -> # there might be a better way to specify this
  dirichlet_prior_ordbeta_stanvar

stanvars <- stanvars + dirichlet_prior_stanvar + dirichlet_prior_ordbeta_stanvar

# Feel free to add any other priors / change the priors on b, 
# which represent regression coefficients on the logit
# scale
  
priors <- set_prior("target += induced_dirichlet_lpdf(thresh | rep_vector(1, 3), 0)", check=FALSE) +
  set_prior("normal(0,5)",class="b") + 
  set_prior("exponential(.1)",class="phi")

# priors <- set_prior("normal(0,5)",class="b") + 
#   prior(constant(0),class="b",coef="Intercept") +
#   prior_string("target += normal_lpdf((cutzero + exp(cutone)) - cutzero|0,3) + cutone",check=F) +
#   set_prior("exponential(.1)",class="phi")

priors_phireg <- set_prior("normal(0,5)",class="b") + 
  set_prior("target += induced_dirichlet_lpdf(thresh | rep_vector(1, 3), 0)", check=FALSE) 
